/*

  Adaptive Huffman compression algorithm implementation
  -----------------------------------------------------

  console mode, single-file RLE/BWT/MFT/Adaptive Huffman compressor

  by yoda in 2002

*/

#define  WIN32_LEAN_AND_MEAN
//#include <windows.h>
#include <stdio.h>
#include <conio.h>
#include <stdlib.h>
//#include <crtdbg.h>
#include "clsFile.h"
#include "MoveToFront.h"
#include "BWT.h"
#include "CRC32.h"
#include "AHuff.h"
#include "StopWatch.h"
#include "RLE.h"

#pragma comment(linker,"/SUBSYSTEM:CONSOLE") // a console, please
#pragma pack(1) // simon says: structures have to be byte aligned

// little optimization
#if !defined(_DEBUG)
#pragma comment(linker,"/FILEALIGN:512 /MERGE:.rdata=.text /MERGE:.data=.text /SECTION:.text,EWR /IGNORE:4078")
#endif

//
// constants
//
#define HC_FLAG_RLE_SKIPPED   (1 << 00)

//
// structs
//

// structures, you'll find in the compressed files
typedef struct _HUFF_COMP_HEADER
{
	WORD                     wSignature;               // HC_SIGNATURE == "HC"
	DWORD                    cbNonShrinked;
	DWORD                    crc;
	DWORD                    flags;                    // HC_FLAG_* (look into the constants section above)

	// SYMBOL_DEF_SHORT[]
} HUFF_COMP_HEADER, *PHUFF_COMP_HEADER;


//
// prototypes
//
BOOL               Compress(PSTR szInFile, PSTR szOutFile);
BOOL               Decompress(PSTR szInFile, PSTR szOutFile);
void               PrintCmdLine();

//
// macros
//
#include "macros.h"

//
// constants
//
#define HC_SIGNATURE          0x4348 // HC -> Huffman compressed

//
// global variables
//

// [...]

//
// root function if compress was specified in the command line
//
BOOL Compress(PSTR szInFile, PSTR szOutFile)
{
	clsFILE            *pfile, *cfOutFile;
	DWORD              dwEncBlockSize, dwOutSize;
	HUFF_COMP_HEADER   hdr;
	BOOL               bRet = FALSE;
	float              fRatio;
	DWORD              cbWorkMem, dwCRC;
	void*              pWorkMem, *pToBeFreed;
	AHuff              huff;
	CStopWatch         clock;
	PBYTE              l_pOutput;

	// pre-init local variables
	pToBeFreed  = NULL;
	l_pOutput   = NULL;

	//
	// map input file
	//
	pfile = new clsFILE();
	if ( !pfile->GetFileHandle(szInFile, F_OPENEXISTING_R) )
	{
		printf("!! Error while opening input file image...\n");
		return FALSE; // ERR
	}
	if ( !pfile->GetFSize() )
	{
		printf("!! 0 byte files aren't allowed as input...\n");
		goto Cleanup; // ERR
	}
	if ( !pfile->MapFile() )
	{
		printf("!! Error while mapping file image...\n");
		goto Cleanup; // ERR
	}
	// already compressed by HuffComp ?
	try
	{
		if ( ((PHUFF_COMP_HEADER)pfile->GetMapPtr())->wSignature == HC_SIGNATURE )
			printf("...input file is already compressed by HuffComp\n");
	}
	catch(...)
	{}

	//
	// paste CRC into header
	//
	ZERO( hdr )
	pWorkMem   = pfile->GetMapPtr();
	cbWorkMem  = pfile->GetMapSize();

	dwCRC = CRC32::Generate( (PBYTE)pWorkMem, cbWorkMem );
	printf( "-> CRC32: 0x%08lX\n", dwCRC );

	//
	// preprocess the input
	//
	printf( "-> Pre-processing input...\n" );
	clock.Start();

	// BWT
	printf( "\tBWT..." );
	if ( BWT::Encode( (PBYTE)pfile->GetMapPtr(), pfile->GetFSize(), pWorkMem, cbWorkMem ) )
		printf( "OK\n" );
	else
		printf( "failed\n" );
	pToBeFreed = pWorkMem; // must be freed later - at the end of this procedure

	// MFT
	printf( "\tMFT..." );
	MoveToFront::Encode( (PBYTE)pWorkMem, cbWorkMem );
	printf( "OK\n" );

	// RLE (eventually shrinks the input buffer !)
	printf( "\tRLE..." );
	hdr.cbNonShrinked = cbWorkMem;

	if ( !RLE::Encode( (PBYTE)pWorkMem, cbWorkMem, cbWorkMem ) )
	{
		hdr.flags |= HC_FLAG_RLE_SKIPPED;
		printf( "skipped (doesn't give a benefit)\n" );
	}
	else
		printf( "OK\n" );

	//
	// encode input data block and write it to the output buffer
	// (out buffer is cleaned by AHuff::Destructor!)
	//
	printf("-> Compressing input data block...\n");
	huff.EncodeData( (PBYTE)pWorkMem, cbWorkMem, (PBYTE &)l_pOutput, dwEncBlockSize );

	clock.Stop();

	//
	// write output file
	//
	printf("-> Writing data to output file...\n");
	cfOutFile = new clsFILE();
	if ( !cfOutFile->GetFileHandle(szOutFile, F_CREATENEW) )
	{
		printf("!! Error while creating output file...\n");
		goto Cleanup; // ERR
	}

	//
	// build & write header
	//
	hdr.wSignature           = HC_SIGNATURE;
	hdr.crc                  = dwCRC;
	cfOutFile->Write( &hdr, sizeof(hdr) );

	//
	// write compressed block
	//
	cfOutFile->Write( l_pOutput, dwEncBlockSize);
	delete cfOutFile;

	//
	// print little summary
	//
	dwOutSize     = dwEncBlockSize + sizeof(HUFF_COMP_HEADER);
	fRatio        = ((float)dwOutSize * 100) / (float)pfile->GetFSize();
	printf( "...Compressed to: %u%%\n", (DWORD)fRatio);
	printf( "...Time elapse: %u min %u s %u ms\n",
		clock.GetTimeMinutes(),
		clock.GetTimeSeconds(),
		clock.GetTimeMSeconds() );
	printf( "-> DONE\n" );
	bRet = TRUE; // return OK

	//
	// cleanup
	//
Cleanup:
	if ( pToBeFreed )
		delete pToBeFreed; // cleanup temporary buffer from BWT encoding
	delete pfile;

	return bRet;
}

//
// root function if decompress was specified in the command line
//

BOOL Decompress(PSTR szInFile, PSTR szOutFile)
{
	clsFILE              *infile, *outfile;
	PHUFF_COMP_HEADER    pHdr;
	DWORD                cbWorkMem;
	void*                pWorkMem = NULL;
	CStopWatch           clock;

	//
	// map input file to memory
	//
	infile = new clsFILE();
	if ( !infile->GetFileHandle(szInFile, F_OPENEXISTING_R) )
	{
		printf("!! Error while opening input file image...\n");
		return FALSE; // ERR
	}
	if ( !infile->GetFSize() )
	{
		printf("!! 0 byte files aren't allowed as input...\n");
		delete infile;
		return FALSE; // ERR
	}
	if ( !infile->MapFile() )
	{
		printf("!! Error while mapping file image...\n");
		delete infile;
		return FALSE; // ERR
	}

	//
	// analyse data block
	//

	// check signature
	void *pInBuff  = infile->GetMapPtr();
	pHdr           = (PHUFF_COMP_HEADER)pInBuff;
	if (pHdr->wSignature != HC_SIGNATURE)
	{
		printf("!! Seems like the input image isn't an output of this tool - bad signature...\n");
		delete infile;
		return FALSE; // ERR
	}

	//
	// calculate some values
	//
	printf("-> Analysing input file image...\n");
	DWORD              dwCompBlockSize = infile->GetFSize() - sizeof(HUFF_COMP_HEADER);
	void*              pCompBlock      = MakePtr( PVOID, pInBuff, sizeof(HUFF_COMP_HEADER) );

	//
	// decode input data block
	//
	printf("-> Decompressing input data block...\n");

	AHuff huff;
	clock.Start();
	huff.DecodeData(
		(PBYTE)pCompBlock, dwCompBlockSize,
		pHdr->cbNonShrinked,
		(PBYTE &)pWorkMem, cbWorkMem );

	//
	// postprocess the output
	//
	printf( "-> Post-processing output buffer...\n" );

	// RLE
	printf( "\tRLE..." );
	if ( pHdr->flags & HC_FLAG_RLE_SKIPPED ) // was original data actually processed by a RLE algo ?
		printf( "skipped\n" );
	else
	{
		RLE::Decode( (PBYTE &)pWorkMem, (DWORD &)cbWorkMem );
		printf( "OK\n" );
	}

	// MFT
	printf( "\tMFT..." );
	MoveToFront::Decode( (PBYTE)pWorkMem, cbWorkMem );
	printf( "OK\n" );

	// BWT
	printf( "\tBWT..." );
	PBYTE pWorkMemCopy = (PBYTE)pWorkMem;
	if ( BWT::Decode( (PBYTE)pWorkMemCopy, pHdr->cbNonShrinked, pWorkMem, cbWorkMem ) )
		printf( "OK\n");
	else
		printf( "failed\n" );
	clock.Stop();

	//
	// validate CRC
	//
	if ( CRC32::Generate( (PBYTE)pWorkMem, cbWorkMem ) == pHdr->crc )
		printf( "-> CRC32 is OK...\n" );
	else
		printf( "-> ! INVALID CRC32 !\n" );

	//
	// write the output buffer to disk
	//
	printf("-> Writing decompressed data to output file...\n");
	outfile = new clsFILE();
	if ( !outfile->GetFileHandle(szOutFile, F_CREATENEW) )
	{
		printf("!! Error while creating output file...\n");
		delete infile;
		delete outfile;
		return FALSE; // ERR
	}
	outfile->Write( pWorkMem, cbWorkMem );

	//
	// print tiny summary
	//
	printf( "...Time elapse: %u min %u s %u ms\n",
		clock.GetTimeMinutes(),
		clock.GetTimeSeconds(),
		clock.GetTimeMSeconds() );
	printf("-> DONE\n");

	//
	// cleanup
	//
	delete infile;
	delete outfile;
//	if ( pWorkMem )
//		delete pWorkMem;

	return TRUE; // OK
}

//
// print command line scheme to STDOUT
//
void PrintCmdLine()
{
	printf("Command line:\n");
	printf("    HuffComp.EXE (input file) (output file) [/d]\n");
	printf("    /d - means decompress input file and save to output file\n");

	return;
}

//
// ENTRYPOINT
//
int main(int argc, char *argv[])
{
	//
	// print title
	//
	printf( "HuffComp 1.3 - Adaptive Huffman compression algorithm implementation\n" );
	printf( "               (including RLE/BWT/MFT data transformators)\n" );
	printf( "by yoda in 2k2\n" );
	printf( "------------------------------------------------------------------\n" );
	printf( "\n");

	//
	// validate command line
	//
	if (argc < 3)
	{
		printf("!! Not enough arguments specified...\n");
		PrintCmdLine();
		return -1;
	}
	if (argc > 4)
	{
		printf("!! Too much arguments specified...\n");
		PrintCmdLine();
		return -1;
	}
	if (argc == 4 && lstrcmpi(argv[3], "/D") != 0)
	{
		PrintCmdLine();
		printf("!! 3rd argument is optional and could only be \"/D\"...\n");
		return -1;
	}

	//
	// compress/decompress
	//
	printf("Processing:  %s...\n", argv[1]);
	printf("Modus:       ");
	__try
	{
		if (argc == 3)
		{
			printf("COMPRESS\n\n");
			Compress(argv[1], argv[2]);
		}
		else
		{
			printf("DECOMPRESS\n\n");
			Decompress(argv[1], argv[2]);
		}
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		printf("!! Exception caught in 'main'...\n");
		return -1;
	}

	return 0;
}